from PIL import Image, ImageDraw, ImageFont
import torch
from torchvision import transforms
from ultralytics import YOLO
import gradio as gr
import numpy as np

# 加载模型和标签
model_detect = YOLO('../model.pt', verbose=False)  # 禁用日志输出
labels = list(model_detect.names.values())  # 动态获取模型的类别标签

# 中英文映射表
en_to_zh_mapping = {
    "crazing": "裂缝",
    "inclusion": "夹杂物",
    "patches": "斑块",
    "pitted_surface": "点蚀表面",
    "rolled-in_scale": "氧化皮",
    "scratches": "划痕"
}

# 将英文类别名称转换为中文
def translate_class_name(class_name):
    return en_to_zh_mapping.get(class_name, class_name)  # 如果找不到对应中文，则返回原英文名

# 辅助函数：解析真实标签文件
def parse_ground_truth(txt_path, img_width, img_height):
    """
    解析真实标签文件，返回边界框列表。
    """
    boxes = []
    try:
        with open(txt_path, 'r') as f:
            for line in f.readlines():
                try:
                    class_id, x_center, y_center, width, height = map(float, line.strip().split())
                    # 将归一化坐标转换为绝对坐标
                    x_min = int((x_center - width / 2) * img_width)
                    y_min = int((y_center - height / 2) * img_height)
                    x_max = int((x_center + width / 2) * img_width)
                    y_max = int((y_center + height / 2) * img_height)
                    boxes.append({
                        "class_id": int(class_id),
                        "bbox": [x_min, y_min, x_max, y_max]
                    })
                except ValueError:
                    print(f"Error parsing line in file {txt_path}: {line}")
    except FileNotFoundError:
        print(f"Label file not found: {txt_path}")
    return boxes

# 辅助函数：计算 IoU
def compute_iou(box1, box2):
    """
    计算两个边界框的 IoU。
    """
    x1, y1, x2, y2 = box1
    x1g, y1g, x2g, y2g = box2

    inter_x1 = max(x1, x1g)
    inter_y1 = max(y1, y1g)
    inter_x2 = min(x2, x2g)
    inter_y2 = min(y2, y2g)

    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
    box1_area = (x2 - x1) * (y2 - y1)
    box2_area = (x2g - x1g) * (y2g - y1g)
    iou = inter_area / (box1_area + box2_area - inter_area)
    return iou

# 辅助函数：计算 AP
def compute_ap(recall, precision):
    """
    计算 AP。
    """
    recall = np.concatenate(([0.0], recall, [1.0]))
    precision = np.concatenate(([0.0], precision, [0.0]))
    for i in range(len(precision) - 1, 0, -1):
        precision[i - 1] = max(precision[i - 1], precision[i])
    indices = np.where(np.diff(recall) > 0)[0]
    ap = np.sum((recall[indices + 1] - recall[indices]) * precision[indices + 1])
    return ap

# 辅助函数：计算 mAP50
def compute_map50(ground_truths, predictions, iou_threshold=0.5):
    """
    计算 mAP50。
    """
    aps = []
    unique_classes = set([gt["class_id"] for gt in ground_truths])

    for cls in unique_classes:
        cls_gts = [gt for gt in ground_truths if gt["class_id"] == cls]
        cls_preds = [pred for pred in predictions if pred["class_id"] == cls]
        cls_preds.sort(key=lambda x: x["confidence"], reverse=True)

        tp = np.zeros(len(cls_preds))
        fp = np.zeros(len(cls_preds))
        matched_gts = set()

        for i, pred in enumerate(cls_preds):
            best_iou = 0
            best_gt_idx = -1
            for j, gt in enumerate(cls_gts):
                iou = compute_iou(pred["bbox"], gt["bbox"])
                if iou > best_iou:
                    best_iou = iou
                    best_gt_idx = j
            if best_iou >= iou_threshold:
                if best_gt_idx not in matched_gts:
                    tp[i] = 1
                    matched_gts.add(best_gt_idx)
                else:
                    fp[i] = 1
            else:
                fp[i] = 1

        tp_cumsum = np.cumsum(tp)
        fp_cumsum = np.cumsum(fp)
        recall = tp_cumsum / len(cls_gts)
        precision = tp_cumsum / (tp_cumsum + fp_cumsum)
        ap = compute_ap(recall, precision)
        aps.append(ap)

    return np.mean(aps)

# 辅助函数：在图片上绘制全局 mAP50 文字
# 辅助函数：在图片上绘制全局 mAP50 文字和图片标号
def draw_map50_and_index_on_image(image, global_map50, image_index):
    """在图片上绘制全局 mAP50 文字和图片标号"""
    draw = ImageDraw.Draw(image)
    try:
        font = ImageFont.truetype("arial.ttf", size=20)  # 加载系统字体
    except IOError:
        font = ImageFont.load_default()  # 如果字体文件不可用，则使用默认字体

    # 绘制全局 mAP50 文字（如果存在）
    if global_map50 is not None:
        map50_text = f"Global mAP50: {global_map50:.2f}"
        draw.text((10, 10), map50_text, fill=(255, 0, 0), font=font)

    # 绘制图片标号
    index_text = f"{image_index + 1}"  # 图片标号从 1 开始
    text_position = (10, image.size[1] - 30)  # 标号位置放在图片底部
    draw.text(text_position, index_text, fill=(0, 0, 255), font=font)  # 蓝色字体
    return image

# 目标检测核心逻辑
# 目标检测核心逻辑
def detect_objects(images, txt_files=None, mode="detect_only"):
    results = []  # 存储每张图片的检测结果
    detection_results_text = []  # 存储每张图片的检测结果文本
    all_ground_truths = []  # 存储所有图片的真实标签
    all_predictions = []  # 存储所有图片的预测结果
    progress = gr.Progress(track_tqdm=True)  # 启用 tqdm 支持

    # 遍历所有图片，解析真实标签并获取预测结果
    for i, image in enumerate(images):
        progress(i / len(images), desc=f"Processing {image.name}")
        try:
            # 加载图像
            img = Image.open(image.name)
            img_width, img_height = img.size

            # 如果是计算 mAP50 模式，解析真实标签
            if mode == "compute_map50":
                if not txt_files or i >= len(txt_files):
                    raise ValueError(f"Missing label file for image {image.name}")
                ground_truths = parse_ground_truth(txt_files[i].name, img_width, img_height)
                all_ground_truths.extend(ground_truths)

            # 获取预测结果
            result = model_detect(image.name)
            predictions = []
            detection_text = f"Image: {image.name}\n"
            for box in result[0].boxes:
                pred = {
                    "class_id": int(box.cls),
                    "bbox": [int(coord) for coord in box.xyxy[0].tolist()],
                    "confidence": float(box.conf)
                }
                predictions.append(pred)
                class_name_en = labels[pred["class_id"]]  # 英文类别名称
                class_name_zh = translate_class_name(class_name_en)  # 转换为中文
                detection_text += f"  缺陷类别: {class_name_zh}, 该缺陷的概率: {pred['confidence']:.2f}, 缺陷的位置: {pred['bbox']}\n"
            detection_results_text.append(detection_text)
            if mode == "compute_map50":
                all_predictions.extend(predictions)

            # 绘制检测结果
            im_array = result[0].plot()
            img_with_boxes = Image.fromarray(im_array[..., ::-1])

            # 在图片上绘制图片标号
            img_with_boxes = draw_map50_and_index_on_image(img_with_boxes, None, i)  # 先不绘制 mAP50
            results.append(img_with_boxes)

        except Exception as e:
            print(f"Error processing image {image.name}: {e}")
            results.append(None)
            detection_results_text.append(f"Error processing image {image.name}: {e}")

    # 如果是计算 mAP50 模式，计算全局 mAP50
    global_map50 = None
    if mode == "compute_map50":
        global_map50 = compute_map50(all_ground_truths, all_predictions)

    # 在每张图片上添加全局 mAP50 文字（如果计算了）
    final_results = []
    for idx, img_with_boxes in enumerate(results):
        if img_with_boxes is not None:
            img_with_boxes = draw_map50_and_index_on_image(img_with_boxes, global_map50, idx)
        final_results.append(img_with_boxes)

    # 将所有检测结果合并为一个字符串
    combined_detection_results = ""
    for idx, detection_text in enumerate(detection_results_text):
        # 替换文件路径为“第几张图片的检测结果”
        formatted_text = f"第 {idx + 1} 张图片的检测结果:\n"
        lines = detection_text.split("\n")[1:]  # 去掉原始文件路径行
        for line in lines:
            if line.strip():  # 跳过空行
                # 替换字段名称
                line = line.replace("Class:", "缺陷类别:")
                line = line.replace("Confidence:", "该缺陷的概率:")
                line = line.replace("BBox:", "缺陷的位置:")
                formatted_text += line + "\n"
        combined_detection_results += formatted_text + "\n\n"

    return final_results, combined_detection_results.strip()

# 创建 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("# 车辆零部件缺陷检测")
    mode = gr.Radio(choices=["仅检测", "计算 mAP50"], value="仅检测", label="选择模式")

    with gr.Row():
        with gr.Column(scale=2):
            input_images = gr.File(file_count="multiple", label="请输入要检测的图片", height=240)
        with gr.Column(scale=1):
            input_txts = gr.File(file_count="multiple", label="请上传对应的标签文件（仅在计算 mAP50 时需要）", height=240, visible=False)
    with gr.Row(height="auto"):
        output_gallery = gr.Gallery(label="检测结果", columns=4, height="auto")
    with gr.Row(height="auto"):
        detection_results_output = gr.Textbox(label="每张图片的检测结果", lines=10)

    def toggle_txt_visibility(mode):
        """根据模式切换标签文件输入框的可见性"""
        return {"visible": mode == "计算 mAP50", "__type__": "update"}

    mode.change(fn=toggle_txt_visibility, inputs=mode, outputs=input_txts)

    button = gr.Button("开始检测", variant="primary")

    def process(mode, images, txt_files):
        if mode == "仅检测":
            return detect_objects(images, mode="detect_only")
        elif mode == "计算 mAP50":
            return detect_objects(images, txt_files, mode="compute_map50")

    button.click(
        fn=process,
        inputs=[mode, input_images, input_txts],
        outputs=[output_gallery, detection_results_output]
    )

# 启动服务
demo.launch(share=True)
